import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal

import tools


class PolicyNet(nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound, num_layers=3):
        """
        Policy network for continuous action space.

        Args:
            state_dim (int): Dimension of the input state.
            hidden_dim (int): Dimension of the hidden layers.
            action_dim (int): Dimension of the action (output) space.
            action_bound (float): Maximum absolute value of the action.
            num_layers (int): Number of hidden layers.
        """
        super(PolicyNet, self).__init__()

        # Input layer
        self.fc1 = nn.Linear(state_dim, hidden_dim)

        # Hidden layers
        self.hidden_layers = nn.ModuleList([
            nn.Linear(hidden_dim, hidden_dim) for _ in range(num_layers - 1)
        ])

        # Output layers for mean and standard deviation
        self.fc_mu = nn.Linear(hidden_dim, action_dim)
        self.fc_std = nn.Linear(hidden_dim, action_dim)

        # Action bound
        self.action_bound = action_bound

    def forward(self, x):
        # Pass through the input layer
        x = F.relu(self.fc1(x))

        # Pass through the hidden layers
        for layer in self.hidden_layers:
            x = F.relu(layer(x))

        # Calculate mean and standard deviation
        mu = self.fc_mu(x)
        std = F.softplus(self.fc_std(x)) #当x非常小的时候它的输出接近0，当x非常大的时候它的输出接近x

        # Define the distribution and sample actions
        dist = Normal(mu, std)
        normal_sample = dist.rsample()  # 通过一个可微分的操作从分布中生成样本，从而支持深度学习中的梯度优化。
        log_prob = dist.log_prob(normal_sample) # 计算这些样本被采样概率的对数。
        action = torch.tanh(normal_sample)

        # Correct the log probability for the tanh transformation
        log_prob = log_prob - torch.log(1 - action.pow(2) + 1e-7)
        """
        这里是修正概率密度,因为 action 是通过 tanh 变换得到的，直接使用 normal_sample 的对数概率密度是不准确的。
        """

        log_prob = log_prob.sum(dim=-1, keepdim=True)  # Sum log probs across action dimensions

        # Scale actions to the action bounds
        action = action * self.action_bound

        return action, log_prob


class QValueNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim, num_hidden_layers=3):
        """
        Q 值网络，支持可控的隐藏层数量。

        Args:
            state_dim (int): 状态的维度。
            hidden_dim (int): 每个隐藏层的神经元数量。
            action_dim (int): 动作的维度。
            num_hidden_layers (int): 隐藏层的数量。
        """
        super(QValueNet, self).__init__()
        self.num_hidden_layers = num_hidden_layers

        # 输入层：连接 state 和 action
        self.fc1 = nn.Linear(state_dim + action_dim, hidden_dim)

        # 隐藏层
        self.hidden_layers = nn.ModuleList([
            nn.Linear(hidden_dim, hidden_dim) for _ in range(num_hidden_layers)
        ])

        # 输出层
        self.fc_out = nn.Linear(hidden_dim, 1)

    def forward(self, x, a):
        # 拼接状态和动作
        cat = torch.cat([x, a], dim=1)

        # 输入层
        x = F.relu(self.fc1(cat))

        # 多个隐藏层
        for hidden_layer in self.hidden_layers:
            x = F.relu(hidden_layer(x))

        # 输出 Q 值
        return self.fc_out(x)


class SAC_Agent:
    """ 处理连续动作的SAC算法 """
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound,
                 actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma,
                 device):
        self.actor = PolicyNet(state_dim, hidden_dim, action_dim, action_bound).to(device)  # 策略网络

        self.critic_1 = QValueNet(state_dim, hidden_dim, action_dim).to(device)  # 第一个Q网络

        self.critic_2 = QValueNet(state_dim, hidden_dim, action_dim).to(device)  # 第二个Q网络

        self.target_critic_1 = QValueNet(state_dim, hidden_dim, action_dim).to(device)  # 第一个目标Q网络


        self.target_critic_2 = QValueNet(state_dim, hidden_dim, action_dim).to(device)  # 第二个目标Q网络


        # 令目标Q网络的初始参数和Q网络一样
        self.target_critic_1.load_state_dict(self.critic_1.state_dict())
        self.target_critic_2.load_state_dict(self.critic_2.state_dict())
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)

        self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(), lr=critic_lr)

        self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(), lr=critic_lr)

        # 使用alpha的log值,可以使训练结果比较稳定
        self.log_alpha = torch.tensor(np.log(0.01), dtype=torch.float)
        self.log_alpha.requires_grad = True  # 可以对alpha求梯度
        self.log_alpha_optimizer = torch.optim.Adam([self.log_alpha],lr=alpha_lr)

        self.target_entropy = target_entropy  # 目标熵的大小
        self.gamma = gamma
        self.tau = tau
        self.device = device
        self.name = "SAC"

    def save_model(self, model_path):
        """保存模型参数"""
        torch.save({
            'actor_state_dict': self.actor.state_dict(),
            'critic_1_state_dict': self.critic_1.state_dict(),
            'critic_2_state_dict': self.critic_2.state_dict(),
            'target_critic_1_state_dict': self.target_critic_1.state_dict(),
            'target_critic_2_state_dict': self.target_critic_2.state_dict(),
        }, model_path)

    def load_model(self, model_path):
        """加载模型参数"""
        checkpoint = torch.load(model_path, map_location=self.device)
        self.actor.load_state_dict(checkpoint['actor_state_dict'])
        self.critic_1.load_state_dict(checkpoint['critic_1_state_dict'])
        self.critic_2.load_state_dict(checkpoint['critic_2_state_dict'])
        self.target_critic_1.load_state_dict(checkpoint['target_critic_1_state_dict'])
        self.target_critic_2.load_state_dict(checkpoint['target_critic_2_state_dict'])


    def take_action(self, state):
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        action = self.actor(state)[0][0]
        return action.tolist()

    def calc_target(self, rewards, next_states, dones):  # 计算目标Q值
        next_actions, log_prob = self.actor(next_states)
        entropy = -log_prob
        q1_value = self.target_critic_1(next_states, next_actions)
        q2_value = self.target_critic_2(next_states, next_actions)
        next_value = torch.min(q1_value,
                               q2_value) + self.log_alpha.exp() * entropy
        td_target = rewards + self.gamma * next_value * (1 - dones)
        return td_target

    def soft_update(self, net, target_net):
        for param_target, param in zip(target_net.parameters(),
                                       net.parameters()):
            param_target.data.copy_(param_target.data * (1.0 - self.tau) + param.data * self.tau)


    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'],dtype=torch.float).to(self.device)

        actions = torch.tensor(transition_dict['actions'],dtype=torch.float).to(self.device)

        rewards = torch.tensor(transition_dict['rewards'],dtype=torch.float).view(-1, 1).to(self.device)

        next_states = torch.tensor(transition_dict['next_states'],dtype=torch.float).to(self.device)

        dones = torch.tensor(transition_dict['dones'],dtype=torch.float).view(-1, 1).to(self.device)

        # 更新两个Q网络
        td_target = self.calc_target(rewards, next_states, dones)
        critic_1_loss = torch.mean(F.mse_loss(self.critic_1(states, actions), td_target.detach()))

        critic_2_loss = torch.mean(F.mse_loss(self.critic_2(states, actions), td_target.detach()))

        self.critic_1_optimizer.zero_grad()
        critic_1_loss.backward()
        self.critic_1_optimizer.step()
        self.critic_2_optimizer.zero_grad()
        critic_2_loss.backward()
        self.critic_2_optimizer.step()

        # 更新策略网络
        new_actions, log_prob = self.actor(states)
        entropy = -log_prob
        q1_value = self.critic_1(states, new_actions)
        q2_value = self.critic_2(states, new_actions)
        actor_loss = torch.mean(-self.log_alpha.exp() * entropy - torch.min(q1_value, q2_value))

        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # 更新alpha值
        alpha_loss = torch.mean((entropy - self.target_entropy).detach() * self.log_alpha.exp())

        self.log_alpha_optimizer.zero_grad()
        alpha_loss.backward()
        self.log_alpha_optimizer.step()

        self.soft_update(self.critic_1, self.target_critic_1)
        self.soft_update(self.critic_2, self.target_critic_2)


class SAC_Agent_PER:
    """使用了优先经验回放"""
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound,
                 actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma, device, buffer_size, batch_size):

        self.actor = PolicyNet(state_dim, hidden_dim, action_dim, action_bound).to(device)  # 策略网络

        self.critic_1 = QValueNet(state_dim, hidden_dim, action_dim).to(device)  # 第一个Q网络

        self.critic_2 = QValueNet(state_dim, hidden_dim, action_dim).to(device)  # 第二个Q网络

        self.target_critic_1 = QValueNet(state_dim, hidden_dim, action_dim).to(device)  # 第一个目标Q网络


        self.target_critic_2 = QValueNet(state_dim, hidden_dim, action_dim).to(device)  # 第二个目标Q网络


        # 令目标Q网络的初始参数和Q网络一样
        self.target_critic_1.load_state_dict(self.critic_1.state_dict())
        self.target_critic_2.load_state_dict(self.critic_2.state_dict())
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)

        self.critic_1_optimizer = torch.optim.Adam(self.critic_1.parameters(), lr=critic_lr)

        self.critic_2_optimizer = torch.optim.Adam(self.critic_2.parameters(), lr=critic_lr)

        # 使用alpha的log值,可以使训练结果比较稳定
        self.log_alpha = torch.tensor(np.log(0.01), dtype=torch.float)
        self.log_alpha.requires_grad = True  # 可以对alpha求梯度
        self.log_alpha_optimizer = torch.optim.Adam([self.log_alpha],lr=alpha_lr)

        self.target_entropy = target_entropy  # 目标熵的大小
        self.gamma = gamma
        self.tau = tau
        self.device = device
        self.memory = tools.ReplayTree(capacity=buffer_size)
        self.batch_size = batch_size
        self.name = 'SAC_PER'

    def save_model(self, model_path):
        """保存模型参数"""
        torch.save({
            'actor_state_dict': self.actor.state_dict(),
            'critic_1_state_dict': self.critic_1.state_dict(),
            'critic_2_state_dict': self.critic_2.state_dict(),
            'target_critic_1_state_dict': self.target_critic_1.state_dict(),
            'target_critic_2_state_dict': self.target_critic_2.state_dict(),
        }, model_path)

    def load_model(self, model_path):
        """加载模型参数"""
        checkpoint = torch.load(model_path, map_location=self.actor.device)
        self.actor.load_state_dict(checkpoint['actor_state_dict'])
        self.critic_1.load_state_dict(checkpoint['critic_1_state_dict'])
        self.critic_2.load_state_dict(checkpoint['critic_2_state_dict'])
        self.target_critic_1.load_state_dict(checkpoint['target_critic_1_state_dict'])
        self.target_critic_2.load_state_dict(checkpoint['target_critic_2_state_dict'])


    def take_action(self, state):
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        action = self.actor(state)[0][0]
        return action.tolist()

    def calc_target(self, rewards, next_states, dones):  # 计算目标Q值
        next_actions, log_prob = self.actor(next_states)
        entropy = -log_prob
        q1_value = self.target_critic_1(next_states, next_actions)
        q2_value = self.target_critic_2(next_states, next_actions)
        next_value = torch.min(q1_value,
                               q2_value) + self.log_alpha.exp() * entropy
        td_target = rewards + self.gamma * next_value * (1 - dones)
        return td_target

    def soft_update(self, net, target_net):
        for param_target, param in zip(target_net.parameters(),
                                       net.parameters()):
            param_target.data.copy_(param_target.data * (1.0 - self.tau) + param.data * self.tau)


    def update(self):
        (states, actions, rewards, next_states, dones), idxs_batch, is_weights_batch = self.memory.sample(
            self.batch_size)
        # 从优先经验回放当中采样数据进行更新。

        states = torch.tensor(states, dtype=torch.float).to(self.device)

        actions = torch.tensor(actions, dtype=torch.float).to(self.device)

        rewards = torch.tensor(rewards, dtype=torch.float).view(-1, 1).to(self.device)

        next_states = torch.tensor(next_states, dtype=torch.float).to(self.device)

        dones = torch.tensor(dones, dtype=torch.float).view(-1, 1).to(self.device)

        is_weights = torch.tensor(is_weights_batch, dtype=torch.float).view(-1, 1).to(self.device)

        # 计算目标 Q 值
        td_target = self.calc_target(rewards, next_states, dones)

        # 计算 Q 网络的损失并加权
        critic_1_td_error = self.critic_1(states, actions) - td_target.detach()
        critic_1_loss = torch.mean(is_weights * critic_1_td_error.pow(2))  # 加权 MSE 损失

        critic_2_td_error = self.critic_2(states, actions) - td_target.detach()
        critic_2_loss = torch.mean(is_weights * critic_2_td_error.pow(2))  # 加权 MSE 损失

        # 更新 Q 网络
        self.critic_1_optimizer.zero_grad()
        critic_1_loss.backward()
        self.critic_1_optimizer.step()

        self.critic_2_optimizer.zero_grad()
        critic_2_loss.backward()
        self.critic_2_optimizer.step()

        # 更新优先级树中的优先级
        td_errors = torch.abs(critic_1_td_error + critic_2_td_error) / 2.0
        self.memory.batch_update(idxs_batch, td_errors.cpu().detach().numpy())

        # 更新策略网络
        new_actions, log_prob = self.actor(states)
        entropy = -log_prob
        q1_value = self.critic_1(states, new_actions)
        q2_value = self.critic_2(states, new_actions)
        actor_loss = torch.mean(-self.log_alpha.exp() * entropy - torch.min(q1_value, q2_value))

        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        # 更新 alpha
        alpha_loss = torch.mean((entropy - self.target_entropy).detach() * self.log_alpha.exp())
        self.log_alpha_optimizer.zero_grad()
        alpha_loss.backward()
        self.log_alpha_optimizer.step()

        # 软更新目标 Q 网络
        self.soft_update(self.critic_1, self.target_critic_1)
        self.soft_update(self.critic_2, self.target_critic_2)


if __name__ == '__main__':
    data = torch.rand(size=(2,20))
    policy = PolicyNet(state_dim=20,hidden_dim=20,action_dim=17,action_bound=0.4)

    r= policy(data)
    print(r,'\n',r)